package com.xjgz.aio.socket.enhance;

import java.io.IOException;
import java.nio.channels.AsynchronousChannelGroup;
import java.nio.channels.SelectionKey;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.nio.channels.Selector;
import java.nio.channels.spi.AsynchronousChannelProvider;

class EnhanceAsynchronousChannelGroup extends AsynchronousChannelGroup {

    /**
     * 递归回调次数上限
     */
    public static final int MAX_INVOKER = 8;

    /**
     * 读回调处理线程池
     */
    private final ExecutorService readExecutorService;

    /**
     * 写回调线程池
     */
    private final ExecutorService commonExecutorService;

    /**
     * write worker
     */
    final Worker commnoWorker;

    /**
     * read workers
     */
    private final Worker[] readWorkers;

    final Worker writeWorker;

    private final AtomicInteger readIndex = new AtomicInteger(0);

    /**
     * group运行状态
     */
    boolean running = true;

    /**
     * 忽略Selection.OP_ACCEPT，该情况下不需要资源回收
     */
    private static final Consumer<SelectionKey> shutdownCallback = selectionKey -> {
        int interestOps = selectionKey.interestOps();
        selectionKey.cancel();
        if ((interestOps & SelectionKey.OP_CONNECT) > 0) {
            Runnable runnable = (Runnable) selectionKey.attachment();
            runnable.run();
        }
        if ((interestOps & SelectionKey.OP_READ) > 0) {
            EnhanceAsynchronousSocketChannel asynchronousSocketChannel =
                (EnhanceAsynchronousSocketChannel) selectionKey.attachment();
            removeOps(selectionKey, SelectionKey.OP_READ);
            asynchronousSocketChannel.doRead(true);
        }
        if ((interestOps & SelectionKey.OP_WRITE) > 0) {
            EnhanceAsynchronousSocketChannel asynchronousSocketChannel =
                (EnhanceAsynchronousSocketChannel) selectionKey.attachment();
            removeOps(selectionKey, SelectionKey.OP_WRITE);
            asynchronousSocketChannel.doWrite();
        }
    };

    private final Consumer<SelectionKey> readConsumer =
        selectionKey -> ((EnhanceAsynchronousSocketChannel) selectionKey.attachment()).doRead(true);

    /**
     * 初始化异步通道组
     * 1.创建读事件处理线程池
     * 2.创建写事件处理线程，使用单据Worker处理写操作
     * 3.创建通用事件处理线程，负责处理accept,connect等操作
     * 4.所有Worker线程都采用Selector机制实现事件监听
     * 
     * @param provider            异步通道提供者
     * @param readExecutorService 读事件处理线程池
     * @param nThreads            读事件处理线程数
     */
    protected EnhanceAsynchronousChannelGroup(
        AsynchronousChannelProvider provider,
        ExecutorService readExecutorService,
        int nThreads
    ) throws IOException {
        super(provider);
        this.readExecutorService = readExecutorService;
        this.readWorkers = new Worker[nThreads];
        for (int i = 0; i < nThreads; i++) {
            this.readWorkers[i] = new Worker(Selector.open(), readConsumer);
            this.readExecutorService.execute(readWorkers[i]);
        }

        writeWorker = new Worker(Selector.open(), selectionKey -> {
            EnhanceAsynchronousSocketChannel asyncchronousSocketChannel =
                (EnhanceAsynchronousSocketChannel) selectionKey.attachment();

            try {
                if (selectionKey.isValid()) {
                    selectionKey.interestOps(selectionKey.interestOps() & ~SelectionKey.OP_WRITE);
                }
            } catch (Throwable e) {
                System.err.println("remove wirte ops error");
                e.printStackTrace();
            } finally {
                while (asyncchronousSocketChannel.doWrite());
            }
        });

        commnoWorker = new Worker(Selector.open(), selectionKey -> {
            if (selectionKey.isAcceptable()) {
                EnhanceAsynchronousServerSocketChannel serverSocketChannel =
                    (EnhanceAsynchronousServerSocketChannel) selectionKey.attachment();
                serverSocketChannel.doAccept();
            } else if (selectionKey.isConnectable()) {
                Runnable runnable = (Runnable) selectionKey.attachment();
                runnable.run();
            } else if (selectionKey.isReadable()) {
                EnhanceAsynchronousSocketChannel asyncchronousSocketChannel =
                    (EnhanceAsynchronousSocketChannel) selectionKey.attachment();
                removeOps(selectionKey, SelectionKey.OP_READ);
                asyncchronousSocketChannel.doRead(true);
            } else {
                throw new IllegalStateException(
                    "unexpect callback, key valid: " + selectionKey.isValid() + ", interestOps: " + selectionKey
                        .interestOps()
                );
            }
        });

        commonExecutorService = new ThreadPoolExecutor(
            2,
            2,
            0L,
            TimeUnit.MILLISECONDS,
            new LinkedBlockingDeque<>(),
            r -> new Thread(r, "smart-socket:common-" + EnhanceAsynchronousChannelGroup.this.hashCode())
        );
        commonExecutorService.execute(writeWorker);
        commonExecutorService.execute(commnoWorker);
    }

    /**
     * 移除指定SelectionKey的指定事件
     * 该方法用于取消对特定IO事件的监听，通常在完成某个IO操作后调用
     * 
     * @param selectionKey 待操作的SelectionKey
     * @param opt          需要移除的事件
     */
    public static void removeOps(SelectionKey selectionKey, int opt) {
        if (selectionKey.isValid() && (selectionKey.interestOps() & opt) != 0) {
            selectionKey.interestOps(selectionKey.interestOps() & ~opt);
        }
    }

    /**
     * @return 返回一个可用的读事件处理Worker
     */
    public Worker getReadWorker() {
        return readWorkers[(readIndex.getAndIncrement() & Integer.MAX_VALUE) % readWorkers.length];
    }


    @Override
    public boolean isShutdown() {
        return readExecutorService.isShutdown();
    }


    @Override
    public boolean isTerminated() {
        return readExecutorService.isTerminated();
    }

    @Override
    public void shutdown() {
        running = false;
        commnoWorker.interrupt();
        writeWorker.interrupt();
        for (Worker worker : readWorkers) {
            worker.interrupt();
        }
        readExecutorService.shutdown();
        commonExecutorService.shutdown();
    }


    @Override
    public void shutdownNow() throws IOException {
        shutdown();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        return readExecutorService.awaitTermination(timeout, unit);
    }

    public static void interestOps(Worker worker, SelectionKey selectionKey, int opt) {
        // 忽略重复注册
        if ((selectionKey.interestOps() & opt) != 0) {
            return;
        }
        // 添加事件
        selectionKey.interestOps(selectionKey.interestOps() | opt);
        if (worker.workThread != Thread.currentThread()) {
            selectionKey.selector().wakeup();
        }
    }

    class Worker implements Runnable {

        /**
         * 当前Worker绑定的Selector
         */
        final Selector selector;

        /**
         * 事件处理回调函数，针对不同的Worker（读/写/通用）有不同的实现
         */
        private final Consumer<SelectionKey> consumer;

        /**
         * 事件注册队列，用于存放待注册的Selector事件
         */
        private final ConcurrentLinkedQueue<Consumer<Selector>> consumers = new ConcurrentLinkedQueue<>();

        /**
         * Worker线程引用，用于线程管理和中断操作
         */
        private Thread workThread;

        public Worker(Selector selector, Consumer<SelectionKey> consumer) {
            this.selector = selector;
            this.consumer = consumer;
        }

        /**
         * 注册新的Selector事件。该方法支持异步注册事件，通过以下机制实现：
         * 1.将注册事件添加到队列中
         * 2.调用Selector.wakeup()方法，唤醒Selector处理新事件
         * 3.确保线程安全的事件注册
         * 
         * @param consumer
         */
        final void addRegister(Consumer<Selector> consumer) {
            consumers.offer(consumer);
            selector.wakeup();
        }

        EnhanceAsynchronousChannelGroup group() {
            return EnhanceAsynchronousChannelGroup.this;
        }

        void interrupt() {
            if (workThread != null) {
                workThread.interrupt();
            }
        }

        @Override
        public void run() {
            workThread = Thread.currentThread();
            Set<SelectionKey> keySet = selector.selectedKeys();
            try {
                while (running) {
                    // 处理待注册的事件
                    Consumer<Selector> selectConsumer;
                    while ((selectConsumer = consumers.poll()) != null) {
                        selectConsumer.accept(selector);
                    }
                    // 阻塞等待IO事件
                    selector.select();
                    // 处理已就绪的IO事件
                    for (SelectionKey key : keySet) {
                        consumer.accept(key);
                    }
                    keySet.clear();
                }
                selector.keys().stream().filter(SelectionKey::isValid).forEach(key -> {
                    try {
                        shutdownCallback.accept(key);
                    } catch (Throwable e) {
                        e.printStackTrace();
                    }
                });
            } catch (Throwable e) {
                if (running) {
                    System.err.println("worker thread error");
                }
                e.printStackTrace();
            } finally {
                try {
                    selector.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }
}
